{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "1f1cea65-1d32-4016-a099-87a94fdfb910",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "# Build an LlamaIndex Agent with UnityCatalog functions\n",
    "In this tutuorial, we'll be covering the steps to use both custom and Databricks-provided AI functions within a LlamaIndex agent. \n",
    "We'll be looking at the tracing integrations with MLflow, as well as utilizing MLflow's models-from-code functionality to simplify the logging, registration, and deployment of our Agent that will use UnityCatalog functions as agent tools. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a78fe7c6-6bdc-4eff-9596-b839e38cdd33",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "%pip install -Uqqq unitycatalog-llamaindex[databricks] openai llama_index mlflow\n",
    "\n",
    "dbutils.library.restartPython()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "3a4d45fa-5fb5-4685-82fe-3b4b515adae6",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import base64\n",
    "import os\n",
    "\n",
    "from databricks.sdk import WorkspaceClient\n",
    "\n",
    "workspace_client = WorkspaceClient()\n",
    "\n",
    "secret_scope = \"ben_wilson\"  # Change me!\n",
    "\n",
    "# Run this if you don't have the API key set to your secrets scope yet\n",
    "\n",
    "# if secret_scope not in [scope.name for scope in workspace_client.secrets.list_scopes()]:\n",
    "#     workspace_client.secrets.create_scope(secret_scope)\n",
    "\n",
    "# my_secret = \"<your API key, temporarily>\"\n",
    "\n",
    "# workspace_client.secrets.put_secret(scope=secret_scope, key=\"openai_api_key\", string_value=my_secret)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "932b7305-cc42-4ed8-af3d-0737c3531b35",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "os.environ[\"OPENAI_API_KEY\"] = base64.b64decode(\n",
    "    workspace_client.secrets.get_secret(scope=secret_scope, key=\"openai_api_key\").value\n",
    ").decode()\n",
    "\n",
    "assert \"OPENAI_API_KEY\" in os.environ, (\n",
    "    \"Please set the OPENAI_API_KEY environment variable to your OpenAI API key\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define constants and a UnityCatalog client instance\n",
    "\n",
    "The constants defined below will be used throughout this tutorial with the marked exception in our agent logging script definition. \n",
    "For understanding context about the requirements for saving an MLflow model using the [models from code](https://mlflow.org/docs/latest/models.html#models-from-code) feature, carefully read the section that\n",
    "defines the script that we'll be saving as a model definition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "afd2b53d-7d82-4e51-91dc-c1939181896c",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from unitycatalog.ai.core.databricks import DatabricksFunctionClient\n",
    "\n",
    "CATALOG = \"ben_wilson\"  # Change me!\n",
    "SCHEMA = \"uc_func\"  # Change me if you want\n",
    "\n",
    "client = DatabricksFunctionClient()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "7b7f7c02-25fb-4f17-b52b-d18a4e55d86a",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Define UC functions\n",
    "\n",
    "In the next several sections, we will be defining 4 distinct UnityCatalog functions.\n",
    "\n",
    "- **execute_python_code**: A hand-crafted function using the `create_python_function` API to register custom functionality from a python function\n",
    "- **ask ai function**: A function that interfaces with the Databricks AI functions services, providing a tool interface for Agents that does not require custom crafting. \n",
    "- **summarization function**: A function that interfaces with the Databricks AI function for summarizing text.\n",
    "- **translation function**: A function that interfaces with the Databricks AI function for providing English <-> Spanish translation without any additional configuration needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "b564c6aa-a9dd-484d-b70f-03f042c52f3a",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionExecutionResult(error=None, format='SCALAR', value='2\\n', truncated=None)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def execute_python_code(code: str) -> str:\n",
    "    \"\"\"\n",
    "    Executes the given python code and returns its stdout.\n",
    "    Remember the code should print the final result to stdout.\n",
    "\n",
    "    Args:\n",
    "      code: Python code to execute. Remember to print the final result to stdout.\n",
    "    \"\"\"\n",
    "    import sys\n",
    "    from io import StringIO\n",
    "\n",
    "    stdout = StringIO()\n",
    "    sys.stdout = stdout\n",
    "    exec(code)\n",
    "    return stdout.getvalue()\n",
    "\n",
    "\n",
    "function_info = client.create_python_function(\n",
    "    func=execute_python_code, catalog=CATALOG, schema=SCHEMA, replace=True\n",
    ")\n",
    "python_execution_function_name = function_info.full_name\n",
    "\n",
    "# test execution\n",
    "client.execute_function(python_execution_function_name, {\"code\": \"print(1+1)\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "93802f13-344b-41f7-ad88-edb2c40d09b9",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Unity Catalog is a feature in Databricks that allows you to manage metadata, such as tables, views, and databases, across multiple workspaces and clouds in a single, unified catalog. It provides a centralized way to manage and govern your data assets, making it easier to discover, understand, and use your data.\\n\\nWith Unity Catalog, you can:\\n\\n1. **Unify metadata management**: Manage metadata across multiple workspaces, clouds, and regions in a single catalog.\\n2. **Simplify data discovery**: Provide a single source of truth for data assets, making it easier for users to find and understand the data they need.\\n3. **Improve data governance**: Apply fine-grained access controls and permissions to ensure that sensitive data is protected and only accessible to authorized users.\\n4. **Enhance collaboration**: Enable multiple teams and users to work together on data projects, while maintaining control and visibility over data assets.\\n\\nUnity Catalog provides a range of features, including:\\n\\n1. **Metadata management**: Create, manage, and version metadata for tables, views, and databases.\\n2. **Data lineage**: Track the origin and movement of data across the organization.\\n3. **Access control**: Apply fine-grained permissions and access controls to data assets.\\n4. **Data discovery**: Provide a search interface for users to find and understand data assets.\\n5. **Integration with Databricks**: Seamlessly integrate with Databricks workspaces, allowing users to access and work with data assets directly from the catalog.\\n\\nOverall, Unity Catalog is designed to help organizations manage their data assets more effectively, improve collaboration and governance, and drive business value from their data.'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ask_ai_function_name = f\"{CATALOG}.{SCHEMA}.ask_ai\"\n",
    "sql_body = f\"\"\"CREATE OR REPLACE FUNCTION {ask_ai_function_name}(question STRING COMMENT 'A question to submit for confirmation by another language model')\n",
    "RETURNS STRING\n",
    "COMMENT 'answer the question by submitting it to the Meta-Llama-3.1-70B-Instruct model'\n",
    "RETURN SELECT ai_gen(question)\n",
    "\"\"\"\n",
    "client.create_function(sql_function_body=sql_body)\n",
    "result = client.execute_function(ask_ai_function_name, {\"question\": \"What is Unity Catalog?\"})\n",
    "result.value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "05030c24-59f0-45da-8dc1-5b1318e62548",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionExecutionResult(error=None, format='SCALAR', value='Unity Catalog: Unified metadata management for Databricks', truncated=None)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summarization_function_name = f\"{CATALOG}.{SCHEMA}.summarize\"\n",
    "sql_body = f\"\"\"CREATE OR REPLACE FUNCTION {summarization_function_name}(text STRING COMMENT 'content that is intended to be summarized', max_words INT COMMENT 'The Maximum number of words to generate within the response. The value must be a non-negative integer. If set to 0, then there is no limit to the length of the response.')\n",
    "RETURNS STRING\n",
    "COMMENT 'Summarize the content and provide a maximum length to the response that will force varying levels of brevity'\n",
    "RETURN SELECT ai_summarize(text, max_words)\n",
    "\"\"\"\n",
    "client.create_function(sql_function_body=sql_body)\n",
    "# test execution\n",
    "client.execute_function(summarization_function_name, {\"text\": result.value, \"max_words\": 20})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a25b0a99-fbe2-41f8-a31e-4d99a42ba2a6",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionExecutionResult(error=None, format='SCALAR', value='¿Qué te gustaría tener para almorzar hoy?</', truncated=None)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "translate_function_name = f\"{CATALOG}.{SCHEMA}.translate\"\n",
    "sql_body = f\"\"\"CREATE OR REPLACE FUNCTION {translate_function_name}(content STRING COMMENT 'Content to translate to the specified language', language STRING COMMENT 'The target language to translate the text to in language shorthand definition. For example, en for English and es for Spanish.')\n",
    "RETURNS STRING\n",
    "COMMENT 'Translate the provided content to the specified target language, currently only a bi-directional translation between  English and Spanish is supported.'\n",
    "RETURN SELECT ai_translate(content, language)\n",
    "\"\"\"\n",
    "client.create_function(sql_function_body=sql_body)\n",
    "# test execution\n",
    "client.execute_function(\n",
    "    translate_function_name,\n",
    "    {\"content\": \"What would you like to have for lunch today?\", \"language\": \"es\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6659aee5-b0a0-462f-89ca-4aa35b628133",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "### Define the locations of the UC functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a9411885-cf00-4d62-9564-bff0edf76471",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "python_execution_function_name = f\"{CATALOG}.{SCHEMA}.execute_python_code\"\n",
    "ask_ai_function_name = f\"{CATALOG}.{SCHEMA}.ask_ai\"\n",
    "summarization_function_name = f\"{CATALOG}.{SCHEMA}.summarize\"\n",
    "translate_function_name = f\"{CATALOG}.{SCHEMA}.translate\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6623cb97-87d0-41c2-b5f2-7828831b3148",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Register Functions as Tools\n",
    "In order for LlamaIndex to 'understand' what the UnityCatalog function interface is, we need to define our functions as tools in the LlamaIndex format. \n",
    "\n",
    "The example below shows how to register all 4 of our functions as LlamaIndex tools, capable of being used within LlamaIndex agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "ca44acf0-eee5-4bfc-ab8e-333170d18cf1",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[UnityCatalogTool(description='Executes the given python code and returns its stdout. Remember the code should print the final result to stdout.', name='ben_wilson__uc_func__execute_python_code', fn_schema=<class 'unitycatalog.ai.core.utils.function_processing_utils.ben_wilson__uc_func__execute_python_code__params'>, return_direct=False),\n",
       " UnityCatalogTool(description='answer the question by submitting it to the Meta-Llama-3.1-70B-Instruct model', name='ben_wilson__uc_func__ask_ai', fn_schema=<class 'unitycatalog.ai.core.utils.function_processing_utils.ben_wilson__uc_func__ask_ai__params'>, return_direct=False),\n",
       " UnityCatalogTool(description='Summarize the content and provide a maximum length to the response that will force varying levels of brevity', name='ben_wilson__uc_func__summarize', fn_schema=<class 'unitycatalog.ai.core.utils.function_processing_utils.ben_wilson__uc_func__summarize__params'>, return_direct=False),\n",
       " UnityCatalogTool(description='Translate the provided content to the specified target language, currently only a bi-directional translation between  English and Spanish is supported.', name='ben_wilson__uc_func__translate', fn_schema=<class 'unitycatalog.ai.core.utils.function_processing_utils.ben_wilson__uc_func__translate__params'>, return_direct=False)]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create tools for LlamaIndex to use\n",
    "\n",
    "from unitycatalog.ai.llama_index.toolkit import UCFunctionToolkit\n",
    "\n",
    "toolkit = UCFunctionToolkit(\n",
    "    client=client,\n",
    "    function_names=[\n",
    "        python_execution_function_name,\n",
    "        ask_ai_function_name,\n",
    "        summarization_function_name,\n",
    "        translate_function_name,\n",
    "    ],\n",
    ")\n",
    "tools = toolkit.tools\n",
    "tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fdeb52c2-b9c7-415e-8a8c-a63fc650798a",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Enable tracing via MLflow\n",
    "\n",
    "[MLflow's LlamaIndex autologging functionality](https://mlflow.org/docs/latest/llms/llama-index/index.html#enable-tracing) allows for traces to be automatically captured with no code modifications needed. We will gain visibility into our agent's internal processes, including the rationalization involved in tool calling, via the [MLflow Tracing UI](https://mlflow.org/docs/latest/llms/tracing/index.html). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "c5ac62bf-22d6-44b4-9a7b-6e04ea5f586a",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import mlflow\n",
    "\n",
    "mlflow.llama_index.autolog()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "b8523e3c-8691-48a7-ba0e-6aefb8de5eff",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Create a LlamaIndex ReActAgent\n",
    "\n",
    "The Agent we're creating is empowering a configured LLM to utilize our UnityCatalog functions as tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "63316479-11cf-4598-b75c-5642675e0c6b",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from llama_index.core.agent import ReActAgent\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "llm = OpenAI()\n",
    "\n",
    "agent = ReActAgent.from_tools(tools, llm=llm, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "a8d10c06-7d20-4ea3-a888-bd6331ec2cbe",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "Submit a chat request to the agent and observe the trace that is recorded from tool calls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1668e328-e343-4c03-bfd7-7428b932ed49",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Running step 1bfe87fe-307c-4646-b2cf-ab2bd9145ad6. Step input: Can you provide a succinct summarization of what Databricks UnityCatalog is in less than 50 words and respond in Spanish?\n",
      "\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
      "Action: ben_wilson__uc_func__translate\n",
      "Action Input: {'content': 'Databricks Unity Catalog is a centralized governance and security layer for data and analytics in the Databricks Lakehouse Platform.', 'language': 'es'}\n",
      "\u001b[0m\u001b[1;3;34mObservation: {\"format\": \"SCALAR\", \"value\": \"El cat\\u00e1logo unificado de Databricks es una capa de gobernanza y seguridad centralizada para datos y an\\u00e1lisis en la plataforma Databricks Lakehouse.\"}\n",
      "\u001b[0m> Running step 4ec6592e-970e-49f5-afe0-0306549ae780. Step input: None\n",
      "\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
      "Answer: El UnityCatalog de Databricks es una plataforma que ofrece activos digitales, herramientas y servicios para el desarrollo de juegos, proporcionando modelos 3D, texturas y complementos para agilizar el proceso de desarrollo e integrar activos de manera eficiente.\n",
      "\u001b[0m"
     ]
    },
    {
     "data": {
      "text/plain": [
       "AgentChatResponse(response='El UnityCatalog de Databricks es una plataforma que ofrece activos digitales, herramientas y servicios para el desarrollo de juegos, proporcionando modelos 3D, texturas y complementos para agilizar el proceso de desarrollo e integrar activos de manera eficiente.', sources=[ToolOutput(content='{\"format\": \"SCALAR\", \"value\": \"El cat\\\\u00e1logo unificado de Databricks es una capa de gobernanza y seguridad centralizada para datos y an\\\\u00e1lisis en la plataforma Databricks Lakehouse.\"}', tool_name='ben_wilson__uc_func__translate', raw_input={'args': (), 'kwargs': {'content': 'Databricks Unity Catalog is a centralized governance and security layer for data and analytics in the Databricks Lakehouse Platform.', 'language': 'es'}}, raw_output='{\"format\": \"SCALAR\", \"value\": \"El cat\\\\u00e1logo unificado de Databricks es una capa de gobernanza y seguridad centralizada para datos y an\\\\u00e1lisis en la plataforma Databricks Lakehouse.\"}', is_error=False)], source_nodes=[], is_dummy_stream=False, metadata=None)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/databricks.mlflow.trace": "\"tr-091608d52741416dbbffdc2fbf589729\"",
      "text/plain": [
       "Trace(request_id=tr-091608d52741416dbbffdc2fbf589729)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "agent.chat(\n",
    "    \"Can you provide a succinct summarization of what Databricks UnityCatalog is in less than 50 words and respond in Spanish?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fa286e9c-7b50-4fe4-ad08-14fdfa5eada4",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Log the Agent to MLflow\n",
    "\n",
    "In order to avoid any potential serialization issues when logging our agent, we're going to use the models-from-code feature in MLflow. This logging approach allows for containment of all dependent logic within a Python script that is then executed when loading the model for inference. \n",
    "\n",
    "We will define the logic in the next cell, utilizing the Jupyter magic command `%%writefile` at the top of the cell to save a local copy of the cell contents as a python script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2a32ceeb-8ebb-44a7-9264-8e7f803a5ebf",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting agent.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile agent.py\n",
    "from mlflow.models import set_model\n",
    "\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core.agent import ReActAgent\n",
    "\n",
    "from unitycatalog.ai.core.client import set_uc_function_client\n",
    "from unitycatalog.ai.core.databricks import DatabricksFunctionClient\n",
    "from unitycatalog.ai.llama_index.toolkit import UCFunctionToolkit\n",
    "\n",
    "\n",
    "# We need to include our constants within the script as it will, upon loading, \n",
    "# run in a separate REPL process.\n",
    "CATALOG = \"ben_wilson\"\n",
    "SCHEMA = \"uc_func\"\n",
    "\n",
    "# The tool calling functionality will need a UnityCatalog functions client to make tool calls\n",
    "# to UnityCatalog with the appropriate authorization.\n",
    "client = DatabricksFunctionClient()\n",
    "\n",
    "# Define our UC functions pathing\n",
    "python_execution_function_name = f\"{CATALOG}.{SCHEMA}.execute_python_code\"\n",
    "ask_ai_function_name = f\"{CATALOG}.{SCHEMA}.ask_ai\"\n",
    "summarization_function_name = f\"{CATALOG}.{SCHEMA}.summarize\"\n",
    "translate_function_name = f\"{CATALOG}.{SCHEMA}.translate\"\n",
    "\n",
    "# Load our UC functions as tools to be used by LlamaIndex\n",
    "toolkit = UCFunctionToolkit(\n",
    "    client=client,\n",
    "    function_names=[\n",
    "        python_execution_function_name,\n",
    "        ask_ai_function_name,\n",
    "        summarization_function_name,\n",
    "        translate_function_name,\n",
    "    ]\n",
    ")\n",
    "tools = toolkit.tools\n",
    "\n",
    "llm = OpenAI()\n",
    "\n",
    "agent = ReActAgent.from_tools(tools, llm=llm, verbose=True)\n",
    "\n",
    "# In order for MLflow to understand what our invocation function is for the model, \n",
    "# we need to specify the interface to the callable as follows:\n",
    "set_model(agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "d73db780-8b74-4b7e-8fc7-7a69cdb730f9",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/11/18 17:01:40 INFO mlflow.llama_index.serialize_objects: API key(s) will be removed from the global Settings object during serialization to protect against key leakage. At inference time, the key(s) must be passed as environment variables.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Running step 9b8079d2-5c8f-41c8-bc0e-f7244a50bc5f. Step input: Where is the largest population of Grey Wolves in the contiguous United States? Please answer in 6 words or less.\n",
      "\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
      "Action: ben_wilson__uc_func__execute_python_code\n",
      "Action Input: {'code': \"print('Yellowstone National Park, Wyoming')\"}\n",
      "\u001b[0m\u001b[1;3;34mObservation: {\"format\": \"SCALAR\", \"value\": \"Yellowstone National Park, Wyoming\\n\"}\n",
      "\u001b[0m> Running step cc53f151-b8f5-4bb7-b04e-f07b74e23b46. Step input: None\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/11/18 17:01:46 INFO mlflow.models.model: Found the following environment variables used during model inference: [OPENAI_API_KEY]. Please check if you need to set them when deploying the model. To disable this message, set environment variable `MLFLOW_RECORD_ENV_VARS_IN_MODEL_LOGGING` to `false`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
      "Answer: Yellowstone National Park, Wyoming.\n",
      "\u001b[0m"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3e6416b1d0b845c59ba36af0bd98944b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading artifacts:   0%|          | 0/14 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Registered model 'ben_wilson.uc_func.llama_index_agent' already exists. Creating a new version of this model...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "66e1203dae98431bbf0f4ca9e7f0d5d1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading artifacts:   0%|          | 0/14 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Created version '3' of model 'ben_wilson.uc_func.llama_index_agent'.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🏃 View run agreeable-moth-983 at: https://e2-dogfood.staging.cloud.databricks.com/ml/experiments/1557189199111224/runs/c99078cdef894beda60dfb74e7576137\n",
      "🧪 View experiment at: https://e2-dogfood.staging.cloud.databricks.com/ml/experiments/1557189199111224\n"
     ]
    }
   ],
   "source": [
    "from mlflow.models import infer_signature\n",
    "\n",
    "input_example = {\n",
    "    \"message\": \"Where is the largest population of Grey Wolves in the contiguous United States? Please answer in 6 words or less.\"\n",
    "}\n",
    "\n",
    "# LLamaIndex's Agent interface is a Pydantic model that is not directly supported by MLflow's signature inference.\n",
    "# To accomodate the output type received from calling the agent, we need to specify what the response will be\n",
    "# from the invocation (in this case, a string).\n",
    "signature = infer_signature(input_example, \"Northern Rocky Mountains, states like Montana.\")\n",
    "\n",
    "# The pip requirements defined in the model logging call are needed solely due to using a pre-release build of\n",
    "# `unitycatalog-ai` and the `unitycatalog-llama_index` packages. When using versions available on PyPI, dependency\n",
    "# inference should work automatically.\n",
    "with mlflow.start_run():\n",
    "    model_info = mlflow.llama_index.log_model(\n",
    "        \"agent.py\",\n",
    "        artifact_path=\"model\",\n",
    "        input_example=input_example,\n",
    "        signature=signature,\n",
    "        pip_requirements=[\n",
    "            \"mlflow\",\n",
    "            \"unitycatalog-ai\",\n",
    "            \"unitycatalog-llamaindex\",\n",
    "            \"llama_index\",\n",
    "            \"openai\",\n",
    "        ],\n",
    "        registered_model_name=\"ben_wilson.uc_func.llama_index_agent\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "52352d19-8053-473f-b5d4-6b4b79ac7ef3",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Validate the ability of our model to be deployed for real-time inference\n",
    "\n",
    "The validation stage here will ensure that the model that we have logged is capable of being deployed to a model serving endpoint. We're validating that the depenendencies are correct, that our input data structure is viable, and that the agent is capable of responding to requests with the configured tools. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "6d0d7ab5-5174-4e29-b2e9-13d418c597d7",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3b65acf8f15c494694491383ace7d417",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/14 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Running step 949962fa-691a-4ee9-a531-91cab4d6eccd. Step input: What is 3**10?\n",
      "\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
      "Action: tool\n",
      "Action Input: {'code': 'print(3**10)'}\n",
      "\u001b[0m\u001b[1;3;34mObservation: Error: No such tool named `tool`.\n",
      "\u001b[0m> Running step 5cf8ab1a-c4f6-43c6-82c4-b038a886ada5. Step input: None\n",
      "\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer\n",
      "Answer: 3 raised to the power of 10 is equal to 59049.\n",
      "\u001b[0m"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'3 raised to the power of 10 is equal to 59049.'"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/databricks.mlflow.trace": "\"tr-7330fa1860554eda8a220b7e2bac93ac\"",
      "text/plain": [
       "Trace(request_id=tr-7330fa1860554eda8a220b7e2bac93ac)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from mlflow.models import convert_input_example_to_serving_input, validate_serving_input\n",
    "\n",
    "serving_input = convert_input_example_to_serving_input({\"message\": \"What is 3**10?\"})\n",
    "validate_serving_input(model_info.model_uri, serving_input=serving_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6be9e50c-ae26-4790-bece-1cf655b18915",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Next Steps\n",
    "From this point, you can safely deploy your model to a serving endpoint. \n",
    "\n",
    "Ensure that the environment variables that were configured at the top of this tutorial (the `OPENAI_API_KEY` and the required Databricks `DATABRICKS_HOST` and `DATABRICKS_TOKEN` variables are set to ensure that the agent can access UnityCatalog for tool calling). "
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "environmentMetadata": {
    "base_environment": "",
    "client": "1"
   },
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 4
   },
   "notebookName": "LlamaIndex UC Agent Databricks functions",
   "widgets": {}
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
